/**
 * 
 */
package org.whuims.bestmatch.tools;

import ireval.RetrievalEvaluator;
import ireval.RetrievalEvaluator.Document;
import ireval.RetrievalEvaluator.Judgment;
import ireval.SetRetrievalComparator;
import ireval.SetRetrievalEvaluator;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.Map;
import java.util.TreeMap;

/**
 * @author ChengQikai
 */
public class SingleFileEval {
	String qrelFilePath = "";
	String rankingFilePath = "";

	public SingleFileEval(String qrelFilePath, String rankingFilePath) {
		super();
		this.qrelFilePath = qrelFilePath;
		this.rankingFilePath = rankingFilePath;
	}

	public double evalNDCG() throws IOException {
		TreeMap<String, ArrayList<Document>> ranking = loadRanking(rankingFilePath);
		TreeMap<String, ArrayList<Judgment>> judgments = loadJudgments(qrelFilePath);
		SetRetrievalEvaluator setEvaluator = create(ranking, judgments);
		// String evalResult = singleEvaluation(setEvaluator, true);
		return setEvaluator.meanNormalizedDiscountedCumulativeGain();
	}

	public String eval() throws IOException {
		TreeMap<String, ArrayList<Document>> ranking = loadRanking(rankingFilePath);
		TreeMap<String, ArrayList<Judgment>> judgments = loadJudgments(qrelFilePath);
		SetRetrievalEvaluator setEvaluator = create(ranking, judgments);
		return singleEvaluation(setEvaluator, true);
	}

	/**
	 * Loads a TREC judgments file.
	 * 
	 * @param filename
	 *            The filename of the judgments file to load.
	 * @return Maps from query numbers to lists of judgments for each query.
	 */

	public static TreeMap<String, ArrayList<Judgment>> loadJudgments(String filename) throws IOException,
			FileNotFoundException {
		// open file
		BufferedReader in = new BufferedReader(new FileReader(filename));
		String line = null;
		TreeMap<String, ArrayList<Judgment>> judgments = new TreeMap<String, ArrayList<Judgment>>();
		String recentQuery = null;
		ArrayList<Judgment> recentJudgments = null;

		while ((line = in.readLine()) != null) {
			// allow for multiple whitespace characters between fields
			String[] fields = line.split("\\s+");

			String number = fields[0];
			String unused = fields[1];
			String docno = fields[2];
			String judgment = fields[3];
			int jVal = 0;
			try {
				jVal = Integer.valueOf(judgment);
			} catch (NumberFormatException e) {
				jVal = (int) Math.round(Double.valueOf(judgment));
			}

			Judgment j = new Judgment(docno, jVal);

			if (recentQuery == null || !recentQuery.equals(number)) {
				if (!judgments.containsKey(number)) {
					judgments.put(number, new ArrayList<Judgment>());
				}

				recentJudgments = judgments.get(number);
				recentQuery = number;
			}

			recentJudgments.add(j);
		}

		in.close();
		return judgments;
	}

	/**
	 * Reads in a TREC ranking file.
	 * 
	 * @param filename
	 *            The filename of the ranking file.
	 * @return A map from query numbers to document ranking lists.
	 */

	public static TreeMap<String, ArrayList<Document>> loadRanking(String filename) throws IOException,
			FileNotFoundException {
		// open file
		BufferedReader in = new BufferedReader(new FileReader(filename));
		String line = null;
		TreeMap<String, ArrayList<Document>> ranking = new TreeMap<String, ArrayList<Document>>();
		ArrayList<Document> recentRanking = null;
		String recentQuery = null;

		while ((line = in.readLine()) != null) {
//			System.out.println(line);
			// allow for multiple whitespace characters between fields
			String[] fields = line.split("\\s+");

			// 1 Q0 WSJ880711-0086 39 -3.05948 Exp

			String number = fields[0];			
			String unused = fields[1];
			String docno = fields[2];
			String rank = fields[3];
			String score = fields[4];
			String runtag = fields[5];
			// lemur can output nan (or NaN)
			double scoreNumber;
			try {
				scoreNumber = Double.valueOf(score);
			} catch (NumberFormatException ex) {
				scoreNumber = 0.0;
			}

			Document document = new Document(docno, Integer.valueOf(rank), scoreNumber);

			if (recentQuery == null || !recentQuery.equals(number)) {
				if (!ranking.containsKey(number)) {
					ranking.put(number, new ArrayList<Document>());
				}

				recentQuery = number;
				recentRanking = ranking.get(number);
			}

			recentRanking.add(document);
		}

		in.close();
		return ranking;
	}

	/**
	 * Creates a SetRetrievalEvaluator from data from loadRanking and
	 * loadJudgments.
	 */

	public static SetRetrievalEvaluator create(TreeMap<String, ArrayList<Document>> allRankings,
			TreeMap<String, ArrayList<Judgment>> allJudgments) {
		// Map query numbers into Integer to get proper sorting.
		TreeMap<String, RetrievalEvaluator> evaluators = new TreeMap<String, RetrievalEvaluator>(
				new java.util.Comparator<String>() {
					public int compare(String a, String b) {
						try {
							Integer a1 = new Integer(a);
							Integer b1 = new Integer(b);
							return a1.compareTo(b1);
						} catch (NumberFormatException e) {
							// not an integer
							return a.compareTo(b);
						}
					}
				});

		for (String query : allRankings.keySet()) {
			ArrayList<Judgment> judgments = allJudgments.get(query);
			ArrayList<Document> ranking = allRankings.get(query);

			/* resort ranking on score, renumber ranks */
			java.util.Collections.sort(ranking, new java.util.Comparator<Document>() {
				public int compare(Document a, Document b) {
					if (a.score < b.score)
						return 1;
					if (a.score == b.score)
						return 0;
					return -1;
				}
			});
			int i = 1;
			for (Document d : ranking) {
				d.rank = i++;
			}

			if (judgments == null || ranking == null) {
				continue;
			}

			RetrievalEvaluator evaluator = new RetrievalEvaluator(query, ranking, judgments);
			evaluators.put(query, evaluator);
		}

		return new SetRetrievalEvaluator(evaluators.values());
	}

	/**
	 * Returns an output string very similar to that of trec_eval.
	 */

	public static String singleQuery(String query, RetrievalEvaluator evaluator) {
		StringWriter s = new StringWriter();
		PrintWriter out = new PrintWriter(s);
		String formatString = "%2$-25s\t%1$5s\t";
		// print trec_eval relational-style output
		// counts
		out.format(formatString + "%3$6d\n", query, "num_ret", evaluator.retrievedDocuments().size());
		out.format(formatString + "%3$6d\n", query, "num_rel", evaluator.relevantDocuments().size());
		out.format(formatString + "%3$6d\n", query, "num_rel_ret", evaluator.relevantRetrievedDocuments().size());

		// aggregate measures
		out.format(formatString + "%3$6.4f\n", query, "map", evaluator.averagePrecision());
		out.format(formatString + "%3$6.4f\n", query, "ndcg", evaluator.normalizedDiscountedCumulativeGain());
		out.format(formatString + "%3$6.4f\n", query, "ndcg15", evaluator.normalizedDiscountedCumulativeGain(15));
		out.format(formatString + "%3$6.4f\n", query, "R-prec", evaluator.rPrecision());
		out.format(formatString + "%3$6.4f\n", query, "bpref", evaluator.binaryPreference());
		out.format(formatString + "%3$6.4f\n", query, "recip_rank", evaluator.reciprocalRank());

		// precision at fixed points
		int[] fixedPoints = RetrievalEvaluator.getFixedPoints();
		double[] vals = evaluator.precisionAtFixedPoints();
		for (int i = 0; i < fixedPoints.length; i++) {
			int point = fixedPoints[i];
			out.format(formatString + "%3$6.4f\n", query, "P" + point, vals[i]);
		}
		double[] precs = evaluator.interpolatedPrecision();
		double prec = 0;
		for (int i = 0; i < precs.length; i++) {
			out.format("at %3$3.2f%2$-18s\t%1$5s\t%4$6.4f\n", query, " ", prec, precs[i]);
			prec += 0.1;
		}
		out.format("\n");
		return s.toString();
	}

	/**
	 * Returns an output string very similar to that of trec_eval.
	 */
	public static String singleEvaluation(SetRetrievalEvaluator setEvaluator, boolean showIndividual) {
		StringWriter s = new StringWriter();
		PrintWriter out = new PrintWriter(s);
		String formatString = "%2$-25s\t%1$5s\t";
		// print trec_eval relational-style output
		if (showIndividual) {
			for (RetrievalEvaluator evaluator : setEvaluator.getEvaluators()) {
				String query = evaluator.queryName();
				out.print(singleQuery(query, evaluator));
			}
		}
		// print summary data
		out.format(formatString + "%3$6d\n", "all", "num_q", setEvaluator.getEvaluators().size());
		out.format(formatString + "%3$6d\n", "all", "num_ret", setEvaluator.numberRetrieved());
		out.format(formatString + "%3$6d\n", "all", "num_rel", setEvaluator.numberRelevant());
		out.format(formatString + "%3$6d\n", "all", "num_rel_ret", setEvaluator.numberRelevantRetrieved());

		out.format(formatString + "%3$6.4f\n", "all", "map", setEvaluator.meanAveragePrecision());
		out.format(formatString + "%3$6.4f\n", "all", "gm_ap", setEvaluator.geometricMeanAveragePrecision());
		out.format(formatString + "%3$6.4f\n", "all", "ndcg", setEvaluator.meanNormalizedDiscountedCumulativeGain());
		out.format(formatString + "%3$6.4f\n", "all", "R-prec", setEvaluator.meanRPrecision());
		out.format(formatString + "%3$6.4f\n", "all", "bpref", setEvaluator.meanBinaryPreference());
		out.format(formatString + "%3$6.4f\n", "all", "recip_rank", setEvaluator.meanReciprocalRank());

		// precision at fixed points
		int[] fixedPoints = SetRetrievalEvaluator.getFixedPoints();
		double[] precs = setEvaluator.precisionAtFixedPoints();

		for (int i = 0; i < fixedPoints.length; i++) {
			int point = fixedPoints[i];
			out.format(formatString + "%3$6.4f\n", "all", "P" + point, precs[i]);
		}
		double prec = 0;
		precs = setEvaluator.interpolatedPrecision();
		for (int i = 0; i < precs.length; i++) {
			out.format("at %3$3.2f%2$-18s\t%1$5s\t%4$6.4f\n", "all", " ", prec, precs[i]);
			prec += 0.1;
		}
		out.format("\n");
		return s.toString();
	}

	public static String comparisonEvaluation(SetRetrievalEvaluator baseline, SetRetrievalEvaluator treatment,
			String baselineName, String treatmentName) {

		StringWriter s = new StringWriter();
		PrintWriter out = new PrintWriter(s);
		String[] metrics = { "averagePrecision", "rPrecision", "ndcg", "bpref", "P5", "P10", "P20" };
		String formatString = "%1$-20s%2$-30s%3$6.4f\n";
		String integerFormatString = "%1$-20s%2$-30s%3$6d\n";
		out.println("Comparing baseline: " + baselineName + " to treatment: " + treatmentName + "\n");
		if (treatment == null)
			return "NOPE";
		for (String metric : metrics) {
			Map<String, Double> baselineMetric = baseline.evaluateAll(metric);
			Map<String, Double> treatmentMetric = treatment.evaluateAll(metric);

			SetRetrievalComparator comparator = new SetRetrievalComparator(baselineMetric, treatmentMetric);

			out.format(formatString, metric, baselineName, comparator.meanBaselineMetric());
			out.format(formatString, metric, treatmentName, comparator.meanTreatmentMetric());

			out.format(integerFormatString, metric, "basebetter", comparator.countBaselineBetter());
			out.format(integerFormatString, metric, "treatbetter", comparator.countTreatmentBetter());
			out.format(integerFormatString, metric, "equal", comparator.countEqual());

			out.format(formatString, metric, "ttest", comparator.pairedTTest());
			out.format(formatString, metric, "randomized", comparator.randomizedTest());
			out.format(formatString, metric, "signtest", comparator.signTest());
		}

		return s.toString();
	}

	public static void usage() {
		System.err.println("ireval: ");
		System.err.println("   There are two ways to use this program.  First, you can evaluate a single ranking: ");
		System.err.println("      java -jar ireval.jar TREC-Ranking-File TREC-Judgments-File");
		System.err.println("   or, you can use it to compare two rankings with statistical tests: ");
		System.err
				.println("      java -jar ireval.jar TREC-Baseline-Ranking-File TREC-Improved-Ranking-File TREC-Judgments-File");
		System.exit(-1);
	}

	public static void main(String[] args) {
		// SingleFileEval eval = new
		// SingleFileEval("resource/qrel/qrels_format.test",
		// "I:\\Resource\\TempGIR\\rel/posRank_sent_stemMode_1_prMode_2_stopWords_false_twScale_1.0_k_100.0.resultrel");
		SingleFileEval eval = new SingleFileEval("resource/qrel/qrels_format.test",
				"F:\\Resource\\TempGir\\rel\\posRank_sent_stemMode_1_prMode_2_stopWords_false_twScale_1.0_k_0.07_lambda_5.0.resultrel");
		try {
			System.out.println(eval.eval());
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
}
